import requests
import json
import time

def run_ollama_http(model, prompt):
    """
    调用 Ollama 本地模型接口（流式输出模式）
    :param model: 模型名称 (如 deepseek-r1:1.5b / qwen2:1.5b)
    :param prompt: 输入提示词
    :return: 模型完整输出
    """
    url = "http://localhost:11434/api/generate"
    payload = {"model": model, "prompt": prompt, "stream": True}

    print(f"\n=== 模型 {model} 输出开始 ===\n")
    response = requests.post(url, json=payload, stream=True)

    result_text = ""
    for line in response.iter_lines():
        if line:
            try:
                data = json.loads(line.decode('utf-8'))
                content = data.get("response", "")
                if content:
                    print(content, end="", flush=True)
                    result_text += content
                if data.get("done"):
                    break
            except Exception:
                continue

    print(f"\n\n=== 模型 {model} 输出结束 ===\n")
    return result_text.strip()


if __name__ == "__main__":
    # 输入文本
    text = "华清宫位于陕西省西安市临潼区骊山脚下，是唐代皇家园林。"

    # 通用 Prompt
    prompt = f"""
请从以下文本中抽取出知识三元组，仅输出结果，不要解释、不带<think>。
输出格式：(实体1, 关系, 实体2) 或 (实体, 属性, 值)
文本：{text}
"""

    # 分别调用 DeepSeek 和 Qwen 模型
    deepseek_output = run_ollama_http("deepseek-r1:1.5b", prompt)
    time.sleep(2)  # 间隔2秒，避免并发冲突
    qwen_output = run_ollama_http("qwen2:1.5b", prompt)

    # 差异统计
    deepseek_lines = [line.strip() for line in deepseek_output.splitlines() if line.strip()]
    qwen_lines = [line.strip() for line in qwen_output.splitlines() if line.strip()]

    print("\n=== 模型结果对比 ===")
    print(f"DeepSeek 输出 {len(deepseek_lines)} 行：{deepseek_lines}")
    print(f"Qwen 输出 {len(qwen_lines)} 行：{qwen_lines}")

    # 简单差异分析
    diff = set(qwen_lines) - set(deepseek_lines)
    if diff:
        print(f"\nQwen 额外识别出的关系：{diff}")
    else:
        print("\n两个模型输出一致 ✅")
